# encoding=utf8
import json
import re
import math
import codecs
from pymongo import MongoClient
from bson import json_util
from urllib.parse import quote_plus
from tqdm import tqdm


user = 'mongouser'
pwd = 'tuanjie123'
host = '118.25.106.14'
port = 27017
url = 'mongodb://%s:%s@%s' % (quote_plus(user), quote_plus(pwd), host)

client = MongoClient(url, port=port)
collection_task_test = client['tagging_test']['task_test']

lines = collection_task_test.find(
    {
        'taskKey': 'sentenceExtract',
        'taskData.segment_results': {'$exists': True},
        "checkData.segment_results": {'$exists': True},
        'canOutput': 1,
        'checkDate': {'$lt': 1640707200000}
    },
    {
        'taskData': 1,
        'checkData': 1,
        'checkDate': 1,
        '_id': 0
    },
    no_cursor_timeout=True
)

# raise Exception('ting')
total = lines.count()
print(total)
# raise Exception("ting")

cut_num = 10000  # 每一万条数据村一个文件
file_index = 97  # 文件标志

save_list = []
success_count = 0
error_count = 0

for idx, line in tqdm(enumerate(lines), total=total):  # lines是动态读取mongo中的数据
    # print(line)
    # raise Exception('ting')
    # try:
    #     cur_news_id = line['newsId']
    #     content = list(collection_content.find({'newsId': cur_news_id}))[0]['content']
    #     line['content'] = re.sub('？，|，？', '？', re.sub('。，|，。', '。', re.sub('，+', '，', re.sub('\s+|</?(.+?)>|,', '，', content.strip()))))
    # except:
    #     error_count += 1
    #     if (idx + 1) % cut_num == 0:  # 达到存储要求
    #         # 将数据写入文件
    #         with codecs.open('./fl_json_file/json_' + str(file_index) + '.json', 'w', encoding='utf-8') as f:
    #             json.dump(save_list, f, ensure_ascii=False, indent=4)
    #         # 清空列表
    #         save_list = []
    #         # 文件标志加一
    #         file_index += 1
    #         last_save = idx
    #
    #     continue
    save_list.append(line)

    if (idx + 1) % cut_num == 0:  # 达到存储要求
        # 将数据写入文件
        with codecs.open('./extract_json_file/json_' + str(file_index) + '.json', 'w', encoding='utf-8') as f:
            json.dump(save_list, f, ensure_ascii=False, indent=4)
        # 清空列表
        save_list = []
        # 文件标志加一
        file_index += 1
        last_save = idx

    # save_list.append(line)
    success_count += 1

# 判断是否有多余数据
if save_list:
    with codecs.open('./extract_json_file/json_' + str(file_index) + '.json', 'w', encoding='utf-8') as f:
        json.dump(save_list, f, ensure_ascii=False, indent=4)

print(success_count)
print(error_count)

